#include "test_tools.h"
#include <mm/early_alloc.h>
#include <mm/page_flags.h>
#include <mm/pages.h>

/// 分配8MB
#define V_MEM_SIZE  (8UL << 20)

#define V_PAGE_ORDER 12

/**
 * @brief 检查是否给定的所有页都被标记为属于伙伴系统。
 */
static inline int __are_all_in_buddy(
    unsigned long physical, 
    unsigned long length
)
{
    struct page_struct *page = physical_to_page((unsigned long)physical);
    if (page == NULL) {
        printf("failed to get page struct. physical: %#lx\n", physical);
        return 0;
    }
    for (int i = 0; i < length; i++) {
        if (!PageBuddy(page + i)) {
            printf("cnt: %d\n", i);
            return 0;
        }
    }
    return 1;
}

/**
 * @brief 检查是否给定的所有页存在被标记为属于伙伴系统。
 */
static inline int __is_any_in_buddy(
    unsigned long physical, 
    unsigned long length
)
{
    struct page_struct *page = physical_to_page((unsigned long)physical);
    if (page == NULL) {
        printf("failed to get page struct. physical: %#lx\n", physical);
        return 0;
    }
    for (int i = 0; i < length; i++) {
        if (PageBuddy(page + i)) {
            return 1;
        }
    }
    return 0;
}

static void early_pages_alloc_test(void)
{
    void *v_mem = aligned_alloc(V_MEM_SIZE, V_MEM_SIZE);
    unsigned long page_start = (VIRT_TO_PHY(v_mem) >> 12);
    unsigned long page_end = page_start + (V_MEM_SIZE >> 12);
    if (early_alloc_init(page_start, page_end, page_start, page_end, 12)) {
        LOG_TEST_FAIL("failed to initialize early_alloc!\n");
    } else {
        LOG_TEST_PASS("early_alloc initilaized!\n");
    }

    if (pages_init(VIRT_TO_PHY(v_mem), VIRT_TO_PHY(v_mem) + V_MEM_SIZE, 0)) {
        LOG_TEST_FAIL("failed to initialize page struct!\n");
    } else {
        LOG_TEST_PASS("page struct initilaized!\n");
    }
    if (set_pages_valid(VIRT_TO_PHY(v_mem), VIRT_TO_PHY(v_mem) + V_MEM_SIZE)) {
        LOG_TEST_FAIL("failed to make page struct valid!\n");
    } else {
        LOG_TEST_PASS("page struct are all valid!\n");
    }

    void *buffer = early_pages_alloc(4, 0);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), (1UL << 4))) {
        LOG_TEST_FAIL("failed to allocate pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate pages!\n");
    }

    buffer = early_pages_alloc(4, 1);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), (1UL << 4))) {
        LOG_TEST_FAIL("failed to allocate aligned pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate aligned pages!\n");
    }

    buffer = early_pages_alloc(8, 0);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), (1UL << 8))) {
        LOG_TEST_FAIL("failed to allocate pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate pages!\n");
    }

    buffer = early_pages_alloc(8, 1);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), (1UL << 8))) {
        LOG_TEST_FAIL("failed to allocate aligned pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate aligned pages!\n");
    }

    buffer = early_pages_alloc(2, 0);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), (1UL << 2))) {
        LOG_TEST_FAIL("failed to allocate pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate pages!\n");
    }

    buffer = early_pages_alloc(2, 1);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), (1UL << 2))) {
        LOG_TEST_FAIL("failed to allocate aligned pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate aligned pages!\n");
    }

    buffer = early_pages_alloc(10, 1);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), (1UL << 10))) {
        LOG_TEST_FAIL("failed to allocate aligned pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate aligned pages!\n");
    }


    early_pages_free(buffer + 4096 * 7, 31);
    if (!__are_all_in_buddy(VIRT_TO_PHY(buffer) + 4096 * 7, 31)) {
        LOG_TEST_FAIL("failed to get pages out of buddy system!\n");
    } else {
        LOG_TEST_PASS("successfully freed pages!\n");
    }
    early_pages_free(buffer + 4096 * 63, 9);
    if (!__are_all_in_buddy(VIRT_TO_PHY(buffer) + 4096 * 63, 9)) {
        LOG_TEST_FAIL("failed to get pages out of buddy system!\n");
    } else {
        LOG_TEST_PASS("successfully freed pages!\n");
    }
    early_pages_free(buffer + 4096 * 128, 8192);
    if (!__are_all_in_buddy(VIRT_TO_PHY(buffer) + 4096 * 128, 1920)) {
        LOG_TEST_FAIL("failed to get pages out of buddy system!\n");
    } else {
        LOG_TEST_PASS("successfully freed pages!\n");
    }
    early_pages_free(v_mem, 64);
    if (!__are_all_in_buddy(VIRT_TO_PHY(v_mem), 64)) {
        LOG_TEST_FAIL("failed to get pages out of buddy system!\n");
    } else {
        LOG_TEST_PASS("successfully freed pages!\n");
    }


    buffer = early_pages_alloc(9, 1);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), (1UL << 9))) {
        LOG_TEST_FAIL("failed to allocate aligned pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate aligned pages!\n");
    }

    buffer = early_pages_alloc(10, 1);
    if (buffer != NULL) {
        LOG_TEST_FAIL("successfully allocate aligned pages, "
            "but it was actually not enough free pages!\n");
        printf("addr: %p, pfn: %ld\n", 
            buffer, 
            (unsigned long)(buffer - v_mem) >> 12
        );
    } else {
        LOG_TEST_PASS("failed to allocate aligned pages "
            "when it was actually not enough free pages!\n");
    }    

    buffer = early_pages_alloc_by_num(1);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), 1)) {
        LOG_TEST_FAIL("failed to allocate pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate pages!\n");
    }

    buffer = early_pages_alloc_by_num(9);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), 9)) {
        LOG_TEST_FAIL("failed to allocate pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate pages!\n");
    }

    buffer = early_pages_alloc_by_num(67);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), 67)) {
        LOG_TEST_FAIL("failed to allocate pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate pages!\n");
    }

    buffer = early_pages_alloc_by_num(141);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), 141)) {
        LOG_TEST_FAIL("failed to allocate pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate pages!\n");
    }

    buffer = early_pages_alloc_by_num(3);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), 3)) {
        LOG_TEST_FAIL("failed to allocate pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate pages!\n");
    }

    buffer = early_pages_alloc_by_num(17);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), 17)) {
        LOG_TEST_FAIL("failed to allocate pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate pages!\n");
    }

    buffer = early_pages_alloc_by_num(67);
    if (buffer == NULL || __is_any_in_buddy(VIRT_TO_PHY(buffer), 67)) {
        LOG_TEST_FAIL("failed to allocate pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate pages!\n");
    }

    free(v_mem);
}

static void early_alloc_test(void)
{
    void *v_mem = aligned_alloc(V_MEM_SIZE, V_MEM_SIZE);
    if (v_mem != NULL) {
        memset(v_mem, 0, V_MEM_SIZE);
    }
    unsigned long page_start = (VIRT_TO_PHY(v_mem) >> 12);
    unsigned long page_end = page_start + (V_MEM_SIZE >> 12);
    if (early_alloc_init(page_start, page_end, page_start, page_end, 12)) {
        LOG_TEST_FAIL("failed to initialize early_alloc!\n");
    } else {
        LOG_TEST_PASS("early_alloc initilaized!\n");
    }

    if (pages_init(VIRT_TO_PHY(v_mem), VIRT_TO_PHY(v_mem) + V_MEM_SIZE, 0)) {
        LOG_TEST_FAIL("failed to initialize page struct!\n");
    } else {
        LOG_TEST_PASS("page struct initilaized!\n");
    }
    if (set_pages_valid(VIRT_TO_PHY(v_mem), VIRT_TO_PHY(v_mem) + V_MEM_SIZE)) {
        LOG_TEST_FAIL("failed to make page struct valid!\n");
    } else {
        LOG_TEST_PASS("page struct are all valid!\n");
    }
    
    void *buffer = early_pages_alloc(8, 1);
    if (buffer == NULL) {
        LOG_TEST_FAIL("failed to allocate aligned pages!\n");
    } else {
        LOG_TEST_PASS("successfully allocate aligned pages!\n");
    }
    memset(buffer, 0xff, 1UL << (12 + 8));
    
    /// 进行分配测试的内存块大小数组。
    size_t alloc_sizes[] = { 
18319, 18718, 20065, 19367, 17607, 19832, 18855, 17939, 18894, 18817, 17377, 17913, 20393, 16792, 19392, 18329,  
16504, 16719, 18018, 17631, 18953, 18314, 17706, 19621, 20247, 19494, 20009, 16842, 20069, 16881, 17061, 16836,  
20174, 16538, 20361, 17590, 20129, 17436, 19202, 20008, 20464, 17382, 19749, 19325, 19752, 19897, 18745, 19786,  
19244, 16752, 18700, 18881, 18993, 17143, 19522, 18331, 19336, 16938, 20370, 16585, 17300, 19457, 19806, 16749,  
16570, 19976, 16673, 20051, 19465, 18658, 20352, 19905, 17636, 16828, 20096, 20343, 17366, 16565, 18762, 17406,  
17665, 16602, 17194, 17767, 18248, 18437, 19616, 19893, 17703, 16897, 17032, 17559, 20175, 17294, 20167, 19489,  
16819, 19651, 20326, 19571, 18719, 17265, 19006, 18003, 17727, 17735, 17993, 19638, 17345, 19080, 17992, 17708,  
20318, 17138, 18407, 18896, 16488, 17291, 18951, 19158, 16599, 18956, 19093, 17974, 17760, 18510, 16428, 16536,  
17478, 16857, 20033, 16795, 20176, 19282, 18499, 18721, 17223, 20258, 20035, 16686, 16491, 17613, 18838, 17965,  
17028, 17367, 17070, 18501, 18462, 18753, 19380, 20336, 17725, 16587, 17648, 16813, 18297, 16452, 17666, 19213,  
18391, 19792, 18315, 19170, 19841, 20181, 16729, 19867, 20328, 19745, 19562, 16604, 19875, 17115, 17813, 17401,  
17371, 18199, 19970, 20450, 18164, 19265, 19444, 20149, 17626, 17044, 18355, 16939, 17166, 19872, 19680, 18144,  
17545, 18232, 18177, 16482, 17082, 16614, 16509, 18585, 19424, 17399, 17675, 18242, 18662, 19140, 19628, 16978,  
16763, 18691, 19092, 18695, 19940, 16901, 18522, 19742, 18412, 17454, 19721, 19254, 19223, 19493, 16975, 20307,  
16822, 19393, 18125, 17637, 17781, 19751, 17538, 16731, 17017, 17959, 17641, 17698, 16675, 19736, 16776, 18548,  
16984, 19981, 16687, 18954, 18353, 16458, 18251, 17024, 18981, 18757, 17048, 19660, 16394, 18371, 16701, 18002,  
17789, 18609, 16808, 19540, 20046, 17145, 17845, 17834, 18652, 16676, 16572, 19008, 17164, 16490, 18254, 20047,  
20044, 16843, 19212, 18862, 19781, 18356, 19125, 18375, 19541, 19309, 17054, 18338, 18780, 17774, 19677, 17657,  
19373, 20406, 18257, 16537, 19518, 17459, 16600, 18146, 19569, 16934, 17528, 16516, 20204, 18318, 20385, 20278,  
17184, 20434, 20384, 18938, 20388, 18240, 17960, 19815, 17742, 17881, 19097, 18021, 19162, 18094, 19509, 18178,  
17821, 16526, 17009, 18055, 19794, 18860, 17552, 18845, 18176, 18968, 20302, 20463, 17536, 18567, 19698, 20034,  
19770, 19418, 18722, 20232, 18387, 17231, 17419, 16434, 19702, 19962, 19707, 19808, 19946, 17212, 19539, 17192,  
16904, 18349, 16757, 16782, 19002, 16706, 20049, 18373, 17716, 18211, 17414, 17848, 17225, 19654, 16913, 
3287, 796, 3174, 1821, 1697, 195, 1733, 650, 2305, 2537, 739, 2613, 1622, 3128, 52, 3636,  
457, 2064, 2946, 64, 2637, 3228, 1225, 517, 958, 1850, 804, 3616, 1542, 2218, 3469, 1979,  
2578, 2940, 2111, 1356, 2176, 1870, 1059, 1025, 1407, 3575, 411, 2202, 2227, 3962, 3983, 3257,  
3997, 1352, 1795, 89, 2157, 1506, 2757, 802, 329, 3011, 1832, 1784, 985, 1885, 3846, 3958,  
1823, 2573, 2149, 1383, 1295, 300, 1626, 3155, 3185, 1277, 2075, 678, 2845, 1829, 3637, 3301,  
1088, 3778, 1044, 2765, 2468, 565, 1590, 2168, 1071, 1237, 192, 1534, 3290, 344, 3118, 378,  
250, 1674, 353, 4028, 583, 826, 2022, 2269, 2444, 3585, 1529, 3009, 3068, 22, 3024, 2025,  
811, 2264, 3775, 2621, 1545, 1122, 2439, 1694, 98, 1648, 3945, 2289, 768, 497, 2789, 3749,  
376, 2875, 2128, 3622, 3386, 1360, 3305, 953, 1654, 470, 3712, 2216, 1230, 1591, 1531, 3604,  
2740, 1912, 2736, 2315, 1268, 3827, 3641, 1464, 2321, 3329, 385, 2337, 431, 446, 205, 
    };
    unsigned long alloc_sizes_len = sizeof(alloc_sizes) / sizeof(size_t);
    buffer = NULL;
    unsigned long cnt = 0;
    unsigned long total_allocated = 0;
    printf("alloc_sizes_len: %lu.\n", alloc_sizes_len);

    /// 挨个进行分配，并进行测试。
    for (cnt = 0; cnt < alloc_sizes_len; cnt++) {
        buffer = early_alloc(alloc_sizes[cnt]);
        if (buffer == NULL) {
            /// 分配失败。
            LOG_TEST_FAIL("cnt: %lu. Allocation failure!\n", cnt);
            break;
        }
        if (buffer < v_mem || buffer > v_mem + V_MEM_SIZE) {
            /// 分配的内存区域首地址越界。
            LOG_TEST_FAIL("cnt: %lu. Start address out of bond!\n", cnt);
            break;
        }
        if (buffer + alloc_sizes[cnt] > v_mem + V_MEM_SIZE) {
            /// 分配的内存区域末尾越界。
            LOG_TEST_FAIL("cnt: %lu. End address out of bond!\n", cnt);
            break;
        }
        if ((unsigned long)buffer & 7) {
            /// 分配的地址没有对齐。
            LOG_TEST_FAIL("cnt: %lu. area is not aligned to 8Bytes!\n", cnt);
            break;
        }

        int to_break = 0;
        for (int i = 0; i < alloc_sizes[cnt]; i++) {
            if (((unsigned char *)buffer)[i] != 0) {
                /// 出现重复分配。
                LOG_TEST_FAIL("cnt: %lu. Multiple allocation!\n", cnt);
                to_break = 1;
                break;
            }
            /// 标记已经被分配过。
            ((unsigned char *)buffer)[i] = 0xff;
        }
        if (to_break) {
            break;
        }
        total_allocated += alloc_sizes[cnt];
    }

    if (cnt == alloc_sizes_len) {
        LOG_TEST_PASS("successfully allocate all blocks!\n");
        printf("total_allocated: %lu\n", total_allocated);
    }
    
    /// 榨干所有内存空间。
    while (early_alloc(8) != NULL);

    /// 空间不足时进行分配的测试
    for (cnt = 0; cnt < alloc_sizes_len; cnt++) {
        buffer = early_alloc(alloc_sizes[cnt]);
        if (buffer != NULL) {
            /// 分配失败。
            LOG_TEST_FAIL("cnt: %lu. Successfully allocated block "
                "but it was already full now!\n", cnt);
            break;
        }
    }

    if (cnt == alloc_sizes_len) {
        LOG_TEST_PASS("no blockk are allocated when it was already full!\n");
    }

    unsigned long pos = 0;
    unsigned long _page_start = 0;
    unsigned long pages_num = 0;
    for_each_early_free_area(pos, _page_start, pages_num) {
        printf("page_start: %lu, pages_num: %lu\n", _page_start, pages_num);
        if (early_pages_alloc(0, 0)) {
            printf("allocation succeeded!\n");
        } else {
            printf("allocation failed!\n");
        }
    }

    free(v_mem);
}


int main(void)
{
    printf("** early_alloc test:\n");
    early_pages_alloc_test();
    early_alloc_test();

    return 0;
}
